#include "utils.h"

extern int gemm(
  const int m, const int n, const int k,
  float *A, const int lda, 
  float *B, const int ldb, 
  float *C, const int ldc,
  void *stream = nullptr);

// JPEGD 128*16, AIPP(VPC) 16*2
int32_t main(int argc, char const *argv[])
{
  std::vector<Case> cases = {};

  cases.push_back({1024, 1024, 1024, 1});

  for(size_t case_id = 0; case_id < cases.size(); ++case_id){
    Case c = cases[case_id];

    const int m = c.m, n = c.n, k = c.k;
    const int loop = c.loop;

    CHECK(aclInit(nullptr))
    int32_t deviceId = 0;
    CHECK(aclrtSetDevice(deviceId))
    aclrtStream stream = nullptr;
    CHECK(aclrtCreateStream(&stream))

    float *h_A = new float[m*k];
    float *h_B = new float[n*k];
    float *d_A = nullptr, *d_B = nullptr, *d_C = nullptr;
    CHECK(aclrtMalloc((void**)&d_A, m*k*sizeof(float), aclrtMemMallocPolicy::ACL_MEM_MALLOC_HUGE_FIRST))
    CHECK(aclrtMalloc((void**)&d_B, n*k*sizeof(float), aclrtMemMallocPolicy::ACL_MEM_MALLOC_HUGE_FIRST))
    CHECK(aclrtMalloc((void**)&d_C, m*n*sizeof(float), aclrtMemMallocPolicy::ACL_MEM_MALLOC_HUGE_FIRST))
    CHECK(aclrtMemcpy(d_A, m*k*sizeof(float), h_A, m*k*sizeof(float), aclrtMemcpyKind::ACL_MEMCPY_HOST_TO_DEVICE))
    CHECK(aclrtMemcpy(d_B, n*k*sizeof(float), h_B, n*k*sizeof(float), aclrtMemcpyKind::ACL_MEMCPY_HOST_TO_DEVICE))
    
    int start = now();
    for(int i = 0; i < loop; ++i){
      gemm(m, n, k, d_A, lda, d_B, ldb, d_C, ldc, stream);
    }
    CHECK(aclrtSynchronizeStream(stream))
    int end = now();

    std::string func_name = "gemm";
    cases[case_id].zhhao_qps = float(loop)/(end-start)*1000;
    dumpbin(d_dst, dst_elem_sz*sizeof(float), func_name + ".bin");
    CHECK(aclrtDestroyStream(stream))
    CHECK(aclrtResetDevice(deviceId))
    CHECK(aclFinalize())
  }
  return 0;
}
